from solvers.solverbase import SolverBase
from solvers.formtest.formtestbase import FormTestBase
from dolfin import *

class BackwardEuler(SolverBase,FormTestBase):
    'Backward-Euler method with extenal forms given.'
    def __init__(self,F,p=1):
        SolverBase.__init__(self,p)
        FormTestBase.__init__(self)
        self.F = F
        self.methodDir = F.methodDir

    def solve(self,problem): # solve problem using given form
        mesh = problem.domain
        n = FacetNormal(mesh)

        p = self.p
        V = FunctionSpace(mesh,'CG',p)
        VV = VectorFunctionSpace(mesh,'CG',p)

        phi_ = interpolate(problem.phi_,V)
        phi = TrialFunction(V)
        varphi = TestFunction(V)

        t = problem.t           
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)


        varphi = varphi + inner(v,nabla_grad(varphi))#stabilization

        a = lhs(self.F(phi,phi_,varphi,v,n,dt,eps=1E-4))
        l = rhs(self.F(phi,phi_,varphi,v,n,dt,eps=1E-4))
  
        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0
       
        A = assemble(a)       
        phi = Function(V)
   
        t = float(dt)
        count = 1

        while t <= T:
            L = assemble(l)  
            solve(A,phi.vector(),L)
            phi_.assign(phi)
            
            self.update(problem,phi_,t,float(dt),count)

            t += float(dt)
            count += 1

        problem.save_data(self.saveDir,str(self.CFL))


